import sys
import os
if os.path.exists("opt"):
    sys.path.append(".")
elif os.path.exists("../opt"):
    sys.path.append("..")

from opt.util.dataset import NeRFDataset 
import argparse
import svox2

# main if
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=str, default="/mnt/data/wangziyi/data/nerf_synthetic/lego")
    parser.add_argument("--split", type=str, default="train")
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--basis_type", type=str, default="sh")
    parser.add_argument("--sh_dim", type=int, default=9)
    parser.add_argument("--basis_reso", type=int, default=16)
    parser.add_argument("--mlp_posenc_size", type=int, default=0)
    parser.add_argument("--mlp_width", type=int, default=16)
    parser.add_argument("--background_nlayers", type=int, default=0)
    parser.add_argument("--background_reso", type=int, default=256)
    parser.add_argument("--reso", type=int, default=256)
    args = parser.parse_args()
    dataset = NeRFDataset(root=args.root, split=args.split)



    grid = svox2.SparseGrid(reso=[args.reso, args.reso, args.reso],
                            center=dataset.scene_center,
                            radius=dataset.scene_radius,
                            use_sphere_bound=dataset.use_sphere_bound,
                            basis_dim=args.sh_dim,
                            use_z_order=True,
                            device=args.device,
                            basis_reso=args.basis_reso,
                            basis_type=svox2.__dict__['BASIS_TYPE_' + args.basis_type.upper()],
                            mlp_posenc_size=args.mlp_posenc_size,
                            mlp_width=args.mlp_width,
                            background_nlayers=args.background_nlayers,
                            background_reso=args.background_reso)

    exit()